#pragma once

#include <iostream>
#include <vector>

namespace nhydata
{
enum Color
{
    RED,BLACK
};

template<class T>
struct RBTreeNode
{   
    using Node = RBTreeNode<T>;
    RBTreeNode(const T& val):
        data_(val),left_(nullptr),right_(nullptr),parent_(nullptr),
        color_(RED)
    {

    }
    Node* left_;
    Node* right_;
    Node* parent_;
    T data_;
    Color color_;
};

template<class T,class Ref,class Ptr>
class __RBTreeIterator
{
public:
    using Node = RBTreeNode<T>;
    using iterator = __RBTreeIterator<T,T&,T*>;
    using Self = __RBTreeIterator<T,Ref,Ptr>;

    __RBTreeIterator() = default;

    __RBTreeIterator(Node* node) : node_(node){}

    __RBTreeIterator(const iterator& it){
        node_ = it.node_;
    }

    Ref operator*(){
        return node_->data_;
    }

    Ptr operator->(){
        return &(node_->data_);
    }

    bool operator!=(const Self& self){
        return node_ != self.node_;
    }

    Self& operator++(){
        if(node_->right_){
            Node* subLeft = node_->right_;
            while(subLeft->left_){
                subLeft = subLeft->left_;
            }
            node_ = subLeft;
        }else {
            Node* cur = node_;
            Node* parent = cur->parent_;
            while(parent && cur == parent->right_){
                cur = parent;
                parent = parent->parent_;
            }
            node_ = parent;
        }
    }

public:
    Node* node_;
};//class __RBTreeIterator end

class KEY
{
public:
    int operator()(const int key){
        return key;
    }
};

template<class K,class T,class KeyOfT>
class RBTree
{
public:
    using Node = RBTreeNode<T>;
    using iterator = __RBTreeIterator<T,T&,T*>;
    using const_iterator = __RBTreeIterator<T,const T&, const T*>;

    ~RBTree(){
        destroy_(root_);
        root_ = nullptr;
    }

    std::pair<iterator,bool> insert(const T& data){
        if(root_ == nullptr){
            root_ = new Node(data);
            root_->color_ = BLACK;
            return std::make_pair(iterator(root_),true);
        }
        //插入
        KeyOfT kot;
        Node* parent = nullptr;
        Node* cur = root_;
        while(cur){
            if(kot(data) < kot(cur->data_)){
                parent = cur;
                cur = cur->left_;
            }else if(kot(data) > kot(cur->data_)) {
                parent = cur;
                cur = cur->right_;
            }else {
                return std::make_pair(iterator(nullptr),false);
            }
        }

        cur = new Node(data);
        Node* newnode = cur;//for return 
        if(kot(data) < kot(parent->data_)){
            parent->left_ = cur;
        }else {
            parent->right_ = cur;
        }
        cur->parent_ = parent;
        //调整
        while(parent && parent->color_ == RED){
            Node* grandfather = parent->parent_;
            if(parent == grandfather->left_){
                Node* uncle = grandfather->right_;
                if(uncle && uncle->color_ == RED){
                    uncle->color_ = BLACK;
                    parent->color_ = BLACK;
                    grandfather->color_ = RED;
                    cur = grandfather;
                    parent = cur->parent_;
                }else { 
                    //uncle 不存在 or uncle是黑色
                    if(cur == parent->left_){
                        rotate_right(grandfather);
                        grandfather->color_ = RED;
                        parent->color_ = BLACK;
                    }else {
                        rotate_left(parent);
                        rotate_right(grandfather);
                        grandfather->color_=RED;
                        cur->color_ = BLACK;
                    }
                    break;
                }
            }else if(parent == grandfather->right_){
                Node* uncle = grandfather->left_;
                if(uncle && uncle->color_ == RED){
                    uncle->color_ = BLACK;
                    parent->color_ = BLACK;
                    grandfather->color_ = RED;
                    cur = grandfather;
                    parent = cur->parent_;
                }else {
                    //uncle不存在 or uncle为黑色
                    if(cur==parent->right_){
                        rotate_left(grandfather);
                        grandfather->color_ = RED;
                        parent->color_ = BLACK;
                    }else {
                        rotate_right(parent);
                        rotate_left(grandfather);
                        grandfather->color_ = RED;
                        cur->color_ = BLACK;
                    }
                    break;
                }

            }
        }//while end

        root_->color_ = BLACK;
        return std::make_pair(iterator(newnode),true);
    }

    Node* find(const K& key){
        KeyOfT kot;
        Node* cur = root_;
        while(cur){
            if(kot(cur->data_) < key){
                cur = cur->right_;
            }else if(kot(cur->data_) < key){
                cur = cur->left_;
            }else {
                return cur;
            }
        }
        return nullptr;
    }
private:
    void rotate_left(Node* parent){
        Node* subR = parent->right_;
        Node* subRL = subR->left_;
        parent->right_ = subRL;
        if(subRL != nullptr)
            subRL->parent_ = parent;
        subR->left_ = parent;
        Node* grandfather =parent->parent_;
        parent->parent_ = subR;
        if(grandfather == nullptr){
            root_ = subR;
            root_->parent_ = nullptr;
        }else {
            if(grandfather->left_ == parent){
                grandfather->left_ = subR;
            }else {
                grandfather->right_ = subR;
            }
            subR->parent_ = grandfather;
        }
    }

    void rotate_right(Node* parent){
        Node* subL = parent->left_;
        Node* subLR = subL->right_;
        parent->left_ = subLR;
        if(subLR != nullptr)
            subLR->parent_ = parent;
        subL->right_ = parent;
        Node* grandfather =parent->parent_;
        parent->parent_ = subL;
        if(grandfather == nullptr){
            root_ = subL;
            root_->parent_ = nullptr;
        }else {
            if(grandfather->left_ == parent){
                grandfather->left_ = subL;
            }else {
                grandfather->right_ = subL;
            }
            subL->parent_ = grandfather;
        }
    }
    
    void InOrder_(Node* root){
        if(root == nullptr) return;
        InOrder_(root->left_);
        std::cout << root->data_ << std::endl;
        InOrder_(root->right_);
    }

    bool is_balance_()
    {
        if(root_ && root_->color_ == RED){
            return false;
        }
        if(!check_(root_,0)) return false;

        return true;
    }

    bool check_(Node* root, int black_num){
        if(root == nullptr){
            std::cout << "path have black_num: " << black_num << std::endl;
            return true;
        }
        if(root->color_ == RED && root->parent_->color_==RED){
            return false;
        } else if(root->color_ == BLACK){
            black_num++;
        }
        return check_(root->left_,black_num) 
            && check_(root->right_,black_num);
    }

    int height_(Node* root){
        if(root == nullptr) return 0;
        int lcount = height_(root->left_);
        int rcount = height_(root->right_);

        return lcount > rcount ? lcount+1 : rcount + 1; 
    }

    void destroy_(Node* root){
        if(root == nullptr) return;
        destroy_(root->left_);
        destroy_(root->right_);

        delete root;
    }
public:

// for test is binary_search_tree
    void in_order(){
        InOrder_(root_);
    }

// for test is balance
    bool is_balance(){
        return is_balance_();
    }
// for test tree height
    int height(){
        return height_(root_);
    }

private:
    Node* root_ = nullptr;
};


}//namespace nhydata end


